import os
import random
import torch
import numpy as np
from PIL import Image, ImageOps, ImageSequence
import folder_paths
import node_helpers


class BatchReferenceLoaderNode:
    """
    批量参考图入库节点
    一次拉入指定数量的本地图，随机打乱，保证每次组合不同
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder_path": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "输入图片文件夹路径"
                }),
                "image_formats": ("STRING", {
                    "default": ".jpg,.jpeg,.png,.bmp,.webp",
                    "multiline": False,
                    "placeholder": "允许的图片格式，用逗号分隔"
                }),
                "count": ("INT", {
                    "default": 9,
                    "min": 1,
                    "max": 50,
                    "step": 1,
                    "display": "number"
                }),
                "seed": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 0xffffffffffffffff,
                    "display": "number"
                }),
                "output_mode": (["separate", "merged"], {
                    "default": "merged"
                }),
            },
            "optional": {
                "merge_layout": (["grid"], {
                    "default": "grid"
                })
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "file_list")
    FUNCTION = "load_batch_images"
    CATEGORY = "AmorForge"

    def load_batch_images(self, folder_path, image_formats, count, seed, output_mode, merge_layout="3x3"):
        # 解析允许的图片格式
        allowed_formats = [fmt.strip().lower() for fmt in image_formats.split(",") if fmt.strip()]
        if not allowed_formats:
            allowed_formats = [".jpg", ".jpeg", ".png", ".bmp", ".webp"]

        # 检查文件夹是否存在
        if not os.path.exists(folder_path):
            raise ValueError(f"文件夹路径不存在: {folder_path}")

        # 获取文件夹中的所有图片文件
        image_files = []
        for file in os.listdir(folder_path):
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in allowed_formats:
                image_files.append(file)

        if not image_files:
            raise ValueError(f"在文件夹 {folder_path} 中没有找到符合条件的图片文件")

        # 设置随机种子并打乱文件列表
        random.seed(seed)
        random.shuffle(image_files)

        # 选择指定数量的图片
        selected_files = image_files[:count] if len(image_files) >= count else image_files * ((count // len(image_files)) + 1)
        selected_files = selected_files[:count]

        # 加载图片
        images = []
        loaded_images = []

        for file in selected_files:
            file_path = os.path.join(folder_path, file)
            try:
                # 使用ComfyUI的标准方式加载图片
                img = node_helpers.pillow(Image.open, file_path)
                img = node_helpers.pillow(ImageOps.exif_transpose, img)

                if img.mode == 'I':
                    img = img.point(lambda i: i * (1 / 255))
                img = img.convert("RGB")

                # 转换为numpy数组并归一化到0-1范围
                img_array = np.array(img).astype(np.float32) / 255.0
                images.append(img_array)
                loaded_images.append(file)
            except Exception as e:
                print(f"警告: 无法加载图片 {file}: {str(e)}")
                continue

        if not images:
            raise ValueError("未能成功加载任何图片")

        # 根据输出模式处理图片
        if output_mode == "separate":
            # 单独输出模式：返回图片批次
            batched_images = torch.from_numpy(np.stack(images))
            file_list = ", ".join(loaded_images)
            return (batched_images, file_list)
        else:
            # 合并输出模式：将图片拼接成一张大图
            merged_image = self.merge_images(images, merge_layout)
            file_list = ", ".join(loaded_images)
            return (merged_image, file_list)

    def merge_images(self, images, layout):
        """将多张图片按指定布局合并成一张大图"""
        if not images:
            raise ValueError("没有图片可以合并")

        # 根据布局确定行列数
        layout_map = {
            "1x1": (1, 1),
            "2x2": (2, 2),
            "3x3": (3, 3),
            "4x4": (4, 4)
        }
        rows, cols = layout_map.get(layout, (3, 3))

        # 计算每张图片的目标尺寸
        total_images = min(len(images), rows * cols)
        grid_images = images[:total_images]

        # 确定统一尺寸（使用第一张图片的尺寸作为基准）
        base_height, base_width = grid_images[0].shape[:2]

        # 创建画布
        canvas_height = base_height * rows
        canvas_width = base_width * cols
        canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.float32)

        # 将图片放置到画布上
        for i, img in enumerate(grid_images):
            row = i // cols
            col = i % cols

            # 调整图片尺寸以适应网格
            resized_img = self.resize_image_to_fit(img, base_height, base_width)

            # 放置到画布上
            y_start = row * base_height
            y_end = y_start + base_height
            x_start = col * base_width
            x_end = x_start + base_width

            canvas[y_start:y_end, x_start:x_end] = resized_img

        # 转换为tensor并添加批次维度
        merged_tensor = torch.from_numpy(canvas)[None,]
        return merged_tensor

    def resize_image_to_fit(self, img, target_height, target_width):
        """调整图片尺寸以适应目标尺寸"""
        img_h, img_w = img.shape[:2]

        # 如果尺寸已经匹配，直接返回
        if img_h == target_height and img_w == target_width:
            return img

        # 调整图片尺寸
        pil_img = Image.fromarray((img * 255).astype(np.uint8))
        resized_pil = pil_img.resize((target_width, target_height), Image.Resampling.LANCZOS)
        resized_np = np.array(resized_pil).astype(np.float32) / 255.0

        return resized_np


# 注册节点
NODE_CLASS_MAPPINGS = {
    "BatchReferenceLoader": BatchReferenceLoaderNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BatchReferenceLoader": "批量参考图加载器"
}